package erasure

import (
	"fmt"

	"github.com/templexxx/reedsolomon"
)

func ErasureEncode(d int, p int, data [][]byte) error {
	if d+p != len(data) {
		return fmt.Errorf("[Erasure] -> data length error: d: %d, p: %d, length data: %d", d, p, len(data))
	}
	for i := 1; i < d; i++ {
		if len(data[0]) != len(data[i]) {
			return fmt.Errorf("[Erasure] -> length of origin data not equal")
		}
	}
	for i := 0; i < p; i++ {
		if len(data[i+d]) != len(data[0]) {
			data[i+d] = make([]byte, len(data[0]))
		}
	}
	ec, err := reedsolomon.New(d, p)
	if err != nil {
		return err
	}
	err = ec.Encode(data)
	if err != nil {
		return err
	}
	return nil
}

func ErasureReconst(d int, p int, data [][]byte, dpHas []int) error {
	if d+p != len(data) {
		return fmt.Errorf("[Erasure] -> data length error: d: %d, p: %d, length data: %d", d, p, len(data))
	}
	for i := 1; i < d+p; i++ {
		if len(data[0]) != len(data[i]) {
			return fmt.Errorf("[Erasure] -> length of data not equal")
		}
	}
	flag := make([]bool, d+p)
	for i := 0; i < len(dpHas); i++ {
		flag[dpHas[i]] = true
	}
	var losts []int
	for i := 0; i < d+p; i++ {
		if !flag[i] {
			losts = append(losts, i)
		}
	}
	ec, err := reedsolomon.New(d, p)
	if err != nil {
		return err
	}
	err = ec.Reconst(data, dpHas, losts)
	if err != nil {
		return err
	}
	return nil
}
